#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;                                     // tensor num for each queue

template<typename TYPE_X1> class KernelDiv {
    using T = TYPE_X1;
public:
    __aicore__ inline KernelDiv() {}
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y,
                                uint32_t totallength, uint32_t tileDataNum) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->totallength = totallength;
        this->tileDataNum = tileDataNum;
        this->alignnum = 32 / sizeof(T);

        x1Gm.SetGlobalBuffer((__gm__ T*)x1, this->totallength + alignnum);
        x2Gm.SetGlobalBuffer((__gm__ T*)x2, this->totallength + alignnum);
        yGm.SetGlobalBuffer((__gm__ T*)y, this->totallength + alignnum);

        pipe.InitBuffer(inQX1, BUFFER_NUM, this->tileDataNum * sizeof(T));
        pipe.InitBuffer(inQX2, BUFFER_NUM, this->tileDataNum * sizeof(T));
        pipe.InitBuffer(outQY, BUFFER_NUM, this->tileDataNum * sizeof(T));
        pipe.InitBuffer(tmp1, this->tileDataNum * sizeof(half));
        pipe.InitBuffer(tmp2, this->tileDataNum * sizeof(half));
        pipe.InitBuffer(tmp3, this->tileDataNum * sizeof(float));
        pipe.InitBuffer(tmp4, this->tileDataNum * sizeof(float));
    }
    __aicore__ inline void Process() {

        int32_t loopCount = this->totallength / this->tileDataNum;
        for (int32_t i = 0; i < loopCount; i++) {
            CopyIn(i,  this->tileDataNum);
            Compute(i, this->tileDataNum);
            CopyOut(i, this->tileDataNum);
        }
        if(this->totallength % this->tileDataNum != 0) {
            uint32_t taillength = this->totallength % this->tileDataNum;
            uint32_t taillengthalign = (taillength + this->alignnum - 1) / this->alignnum * this->alignnum;
            CopyIn(loopCount,  taillengthalign);
            Compute(loopCount, taillength);
            CopyOut(loopCount, taillengthalign);
        }
    }
    __aicore__ inline void InitK2BR(GM_ADDR x1, GM_ADDR x2, GM_ADDR y,
                                uint32_t totallength, uint32_t tileDataNum,
                                uint32_t x1_length, uint32_t x2_length, 
                                int64_t numshapes, int64_t ss[], int64_t sf[]) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->tileDataNum = tileDataNum;
        this->alignnum = 32 / sizeof(T);
        this->numshapes = numshapes;
        this->x1_Length = x1_length;
        this->x2_Length = x2_length;
        for (int i = 0; i < 128; ++i) {
            ((int64_t *)this->shape)[i] = ss[i];
        }
        for(int i = 0; i < 64; ++i) {
            ((int64_t *)this->shapefull)[i] = sf[i];
        }        
        this->totallength = 1;
        for(int i = 0; i < this->numshapes; ++i) {
            this->totallength *= ((int64_t *)this->shapefull)[i];
        }
        x1Gm.SetGlobalBuffer((__gm__ T*)x1, this->totallength + alignnum);
        x2Gm.SetGlobalBuffer((__gm__ T*)x2, this->totallength + alignnum);
        yGm.SetGlobalBuffer((__gm__ T*)y, this->totallength + alignnum);

        pipe.InitBuffer(inQX1, BUFFER_NUM, this->tileDataNum * sizeof(T));
        pipe.InitBuffer(inQX2, BUFFER_NUM, this->tileDataNum * sizeof(T));
        pipe.InitBuffer(outQY, BUFFER_NUM, this->tileDataNum * sizeof(T));
        pipe.InitBuffer(tmp1, this->tileDataNum * sizeof(half));
        pipe.InitBuffer(tmp2, this->tileDataNum * sizeof(half));
        pipe.InitBuffer(tmp3, this->tileDataNum * sizeof(float));
        pipe.InitBuffer(tmp4, this->tileDataNum * sizeof(float));
    }
    __aicore__ inline void ProcessK2BR() {

        int32_t loopCount = this->totallength / this->tileDataNum;
        for (int32_t i = 0; i < loopCount; i++) {
            CopyInK2BR(i,  this->tileDataNum);
            Compute(i, this->tileDataNum);
            CopyOut(i, this->tileDataNum);
        }
        if(this->totallength % this->tileDataNum != 0) {
            uint32_t taillength = this->totallength % this->tileDataNum;
            uint32_t taillengthalign = (taillength + this->alignnum - 1) / this->alignnum * this->alignnum;
            CopyInK2BR(loopCount,  taillengthalign);
            Compute(loopCount, taillength);
            CopyOut(loopCount, taillengthalign);
        }
    }

private:
    __aicore__ inline void CopyIn(int32_t progress, uint32_t length)
    {
        LocalTensor<T> x1Local = inQX1.AllocTensor<T>();
        DataCopy(x1Local, x1Gm[progress * this->tileDataNum], length);
        inQX1.EnQue(x1Local);

        LocalTensor<T> x2Local = inQX2.AllocTensor<T>();
        DataCopy(x2Local, x2Gm[progress * this->tileDataNum], length);
        inQX2.EnQue(x2Local);
    }

    __aicore__ inline void Compute(int32_t progress, uint32_t length)
    {
        // if constexpr (std::is_same_v<TYPE_X, half>) {
        //     Computefp16(progress);
        //     return;
        // }    
        LocalTensor<T> x1Local = inQX1.DeQue<T>();
        LocalTensor<T> x2Local = inQX2.DeQue<T>();
        LocalTensor<T> yLocal = outQY.AllocTensor<T>();
        if constexpr (std::is_same_v<T, int8_t> || std::is_same_v<T, signed char>) {
            auto p1 = tmp1.Get<half>();
            auto p2 = tmp2.Get<half>();
            Cast(p1, x1Local, RoundMode::CAST_NONE, length);
            Cast(p2, x2Local, RoundMode::CAST_NONE, length);
            Div(p1, p1, p2, length);
            Cast(yLocal, p1, RoundMode::CAST_TRUNC, length);
        } else if constexpr (std::is_same_v<T, int32_t>){
            auto p1 = tmp3.Get<float>();
            auto p2 = tmp4.Get<float>();
            Cast(p1, x1Local, RoundMode::CAST_ROUND, length);
            Cast(p2, x2Local, RoundMode::CAST_ROUND, length);
            Div(p1, p1, p2, length);
            Cast(yLocal, p1, RoundMode::CAST_TRUNC, length);
        } else {
            Div(yLocal, x1Local, x2Local, length);
        }

        outQY.EnQue<T>(yLocal);
        inQX1.FreeTensor(x1Local);
        inQX2.FreeTensor(x2Local);
    }
    __aicore__ inline void CopyOut(int32_t progress, uint32_t length)
    {
        LocalTensor<T> yLocal = outQY.DeQue<T>();
        DataCopy(yGm[progress * this->tileDataNum], yLocal, length);
        outQY.FreeTensor(yLocal);
    }
    __aicore__ inline void CopyInK2BR(int32_t progress, uint32_t length)
    {
        LocalTensor<T> x1Local = inQX1.AllocTensor<T>();
        //DataCopy(x1Local, x1Gm[progress * this->tileDataNum], length);
        if(this->x1_Length < this->totallength) {
            BroadCX1(x1Local, progress * this->tileDataNum, length);
        } else {            
            DataCopy(x1Local, x1Gm[progress * this->tileDataNum], length);
        }
        inQX1.EnQue(x1Local);

        LocalTensor<T> x2Local = inQX2.AllocTensor<T>();
        //DataCopy(x2Local, x2Gm[progress * this->tileDataNum], length);
        if(this->x2_Length < this->totallength) {
            BroadCX2(x2Local, progress * this->tileDataNum, length);
        } else {
            DataCopy(x2Local, x2Gm[progress * this->tileDataNum], length);
        }
        inQX2.EnQue(x2Local);
    }
    __aicore__ inline void BroadCX1(LocalTensor<T> &dst, uint32_t offset, uint32_t length) {
        if(this->x1_Length == 1) {
            if constexpr (std::is_same_v<T, int32_t> || std::is_same_v<T, half> || std::is_same_v<T, float>) {
                T tmp = x1Gm.GetValue(0);
                Duplicate(dst, tmp, length);
                return;
            }
        }
        for(uint32_t i = 0; i < length; i++) {
            int idxtmp = 0;
            int istart = i + offset;
            for(int k = 1; k <= this->numshapes; k++) {
                int kpos = 0;
                int krange = 1;
                if(k < this->numshapes) {
                    for(int m = k + 1; m <= this->numshapes; m++) {
                        krange *= shapefull[m - 1];
                    }
                    kpos = istart / krange;
                    istart = istart % krange;
                } else {
                    krange = shapefull[k - 1];
                    kpos = istart % krange;
                }
                //idxtmp += kpos * this->stride[k - 1];
                int krangeB = 1;
                if(shapefull[k - 1] == shape[0][k - 1]) {
                    if(k < this->numshapes) {
                        for(int m = k + 1; m <= this->numshapes; m++) {
                            krangeB *= shape[0][m - 1];
                        }
                        idxtmp += kpos * krangeB;
                    }  else {
                        idxtmp += kpos;
                    }
                }
            }
            T tmp = x1Gm.GetValue(idxtmp);
            dst.SetValue(i, tmp);
        }
    }
    __aicore__ inline void BroadCX2(LocalTensor<T> &dst, uint32_t offset, uint32_t length) {
        if(this->x2_Length == 1) {
            if constexpr (std::is_same_v<T, int32_t> || std::is_same_v<T, half> || std::is_same_v<T, float>) {
                T tmp = x2Gm.GetValue(0);
                Duplicate(dst, tmp, length);
                return;
            }    
        }
        for(uint32_t i = 0; i < length; i++) {
            int idxtmp = 0;
            int istart = i + offset;
            for(int k = 1; k <= this->numshapes; k++) {
                int kpos = 0;
                int krange = 1;
                if(k < this->numshapes) {
                    for(int m = k + 1; m <= this->numshapes; m++) {
                        krange *= shapefull[m - 1];
                    }
                    kpos = istart / krange;
                    istart = istart % krange;
                } else {
                    krange = shapefull[k - 1];
                    kpos = istart % krange;
                }
                //idxtmp += kpos * this->stride[k - 1];
                int krangeB = 1;
                if(shapefull[k - 1] == shape[1][k - 1]) {
                    if(k < this->numshapes) {
                        for(int m = k + 1; m <= this->numshapes; m++) {
                            krangeB *= shape[1][m - 1];
                        }
                        idxtmp += kpos * krangeB;
                    }  else {
                        idxtmp += kpos;
                    }
                }
            }
            T tmp = x2Gm.GetValue(idxtmp);
            dst.SetValue(i, tmp);
        }
    }
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQX1, inQX2;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQY;
    TBuf<QuePosition::VECCALC> tmp1, tmp2, tmp3, tmp4;

    GlobalTensor<T> x1Gm;
    GlobalTensor<T> x2Gm;
    GlobalTensor<T> yGm;
    uint32_t totallength;
    uint32_t tileDataNum;
    uint32_t x1_Length;
    uint32_t x2_Length;
    uint32_t alignnum;
    int64_t shape[2][64];
    int64_t numshapes;
    int64_t shapefull[64];
};

extern "C" __global__ __aicore__ void div(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    KernelDiv<DTYPE_X1> op;
    if (TILING_KEY_IS(1)) {
        op.Init(x1, x2, y, 
                tiling_data.totallength, tiling_data.tileDataNum);  
        op.Process();
    } else if (TILING_KEY_IS(2)) {
        op.InitK2BR(x1, x2, y, 
                tiling_data.totallength, tiling_data.tileDataNum,
                tiling_data.x1_length, tiling_data.x2_length, 
                tiling_data.numshapes, tiling_data.shape, tiling_data.shapefull);  
        op.ProcessK2BR();
    }    
}